#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：yolo_data_preparation.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/1/7 上午9:22
"""

import os
import shutil
import random
import json
from typing import List


class YOLOv5DataPreparer:
    def __init__(self, input_dir: str, output_dir: str, train_ratio: float = 0.8):
        """
        初始化YOLOv5数据准备器。

        :param input_dir: 包含images和labels文件夹的输入目录路径。
        :param output_dir: 输出YOLO格式数据的目录路径。
        :param train_ratio: 训练集和测试集的比例，默认为0.8。
        """
        self.input_dir = os.path.abspath(input_dir)
        self.output_dir = os.path.abspath(output_dir)
        self.train_ratio = train_ratio
        self.images_dir = os.path.join(self.input_dir, 'images')
        self.labels_dir = os.path.join(self.input_dir, 'labels')
        self.train_images_dir = os.path.join(self.output_dir, 'train', 'images')
        self.train_labels_dir = os.path.join(self.output_dir, 'train', 'labels')
        self.test_images_dir = os.path.join(self.output_dir, 'test', 'images')
        self.test_labels_dir = os.path.join(self.output_dir, 'test', 'labels')
        os.makedirs(self.train_images_dir, exist_ok=True)
        os.makedirs(self.train_labels_dir, exist_ok=True)
        os.makedirs(self.test_images_dir, exist_ok=True)
        os.makedirs(self.test_labels_dir, exist_ok=True)

        self.classes = self._get_classes()
        if not self.classes:
            print("警告：未找到任何类别标签。")
        self._write_classes_file()

    def _get_classes(self) -> List[str]:
        """
        从JSON文件中提取所有类别。

        :return: 类别列表。
        """
        classes = set()
        for filename in os.listdir(self.labels_dir):
            if filename.endswith('.json'):
                try:
                    with open(os.path.join(self.labels_dir, filename), 'r') as f:
                        data = json.load(f)
                        for shape in data.get('shapes', []):
                            classes.add(shape['label'])
                except Exception as e:
                    print(f"读取文件 {filename} 时出错: {e}")
        return sorted(list(classes))

    def _write_classes_file(self):
        """
        将类别写入classes.txt文件。
        """
        if self.classes:
            with open(os.path.join(self.output_dir, 'classes.txt'), 'w') as f:
                for cls in self.classes:
                    f.write(cls + '\n')

    def _convert_shape(self, shape: dict, img_width: int, img_height: int) -> List[float]:
        """
        将Labelme的矩形形状转换为YOLO格式。

        :param shape: Labelme形状字典。
        :param img_width: 图像宽度。
        :param img_height: 图像高度。
        :return: YOLO格式的中心点坐标和宽高。
        """
        x1, y1 = shape['points'][0]
        x2, y2 = shape['points'][1]
        x_center = (x1 + x2) / 2 / img_width
        y_center = (y1 + y2) / 2 / img_height
        width = (x2 - x1) / img_width
        height = (y2 - y1) / img_height
        return [x_center, y_center, width, height]

    def _convert_json_to_yolo(self, json_path: str, output_dir: str):
        """
        将单个JSON文件转换为YOLO格式。

        :param json_path: JSON文件路径。
        :param output_dir: 输出目录。
        """
        with open(json_path, 'r') as f:
            data = json.load(f)
            img_width, img_height = data['imageWidth'], data['imageHeight']
            txt_filename = os.path.splitext(os.path.basename(json_path))[0] + '.txt'
            with open(os.path.join(output_dir, txt_filename), 'w') as txt_file:
                for shape in data.get('shapes', []):
                    if shape['shape_type'] == 'rectangle':
                        x_center, y_center, width, height = self._convert_shape(shape, img_width, img_height)
                        class_id = self.classes.index(shape['label'])
                        txt_file.write(f"{class_id} {x_center} {y_center} {width} {height}\n")

    def _copy_image(self, image_path: str, output_dir: str):
        """
        复制图像文件到指定目录。

        :param image_path: 图像文件路径。
        :param output_dir: 输出目录。
        """
        try:
            shutil.copy(image_path, output_dir)
        except FileNotFoundError:
            print(f"图像文件 {image_path} 未找到，跳过复制。")
        except Exception as e:
            print(f"复制图像文件 {image_path} 时出错: {e}")

    def prepare_dataset(self):
        """
        准备数据集，包括转换标签和复制图像。
        """
        json_files = [f for f in os.listdir(self.labels_dir) if f.endswith('.json')]
        random.shuffle(json_files)
        train_size = int(len(json_files) * self.train_ratio)
        train_files = json_files[:train_size]
        test_files = json_files[train_size:]

        for json_file in train_files:
            json_path = os.path.join(self.labels_dir, json_file)
            self._convert_json_to_yolo(json_path, self.train_labels_dir)
            image_filename = os.path.splitext(json_file)[0] + '.jpg'
            image_path = os.path.join(self.images_dir, image_filename)
            self._copy_image(image_path, self.train_images_dir)

        for json_file in test_files:
            json_path = os.path.join(self.labels_dir, json_file)
            self._convert_json_to_yolo(json_path, self.test_labels_dir)
            image_filename = os.path.splitext(json_file)[0] + '.jpg'
            image_path = os.path.join(self.images_dir, image_filename)
            self._copy_image(image_path, self.test_images_dir)

        # 检查目录是否为空
        if not any(os.scandir(self.train_images_dir)):
            print(f"警告：{self.train_images_dir} 目录为空，请检查数据准备过程。")
        if not any(os.scandir(self.train_labels_dir)):
            print(f"警告：{self.train_labels_dir} 目录为空，请检查数据准备过程。")
        if not any(os.scandir(self.test_images_dir)):
            print(f"警告：{self.test_images_dir} 目录为空，请检查数据准备过程。")
        if not any(os.scandir(self.test_labels_dir)):
            print(f"警告：{self.test_labels_dir} 目录为空，请检查数据准备过程。")

        print(f"数据集准备完成。训练集: {len(train_files)}个文件，测试集: {len(test_files)}个文件。")


# 使用示例
if __name__ == "__main__":
    input_dir = 'datasets/HELMET/source_data'
    output_dir = 'datasets/HELMET/yolo_output_dir'
    data_preparer = YOLOv5DataPreparer(input_dir, output_dir)
    data_preparer.prepare_dataset()